import asyncio
import base64
from io import BytesIO
from PIL import ImageGrab
import tkinter as tk
from tkinter import messagebox
import websockets
import threading
import time  # 导入 time 模块

# 全局变量，用于控制WebSocket服务器的状态
server = None
loop = None

async def handle_screenshot(websocket, path=None):
    async for message in websocket:
        try:
            # 解析收到的消息
            x1, y1, x2, y2 = map(int, message.split(','))
            # 记录开始时间
            start_time = time.time()
            
            # 截图 该方法花费大概20~40毫秒
            img = ImageGrab.grab(bbox=(x1, y1, x2, y2))
            
            # 记录结束时间
            end_time = time.time()
            
            # 计算并打印花费的时间（毫秒）
            elapsed_time_ms = (end_time - start_time) * 1000
            print(f"Screenshot took {elapsed_time_ms:.2f} ms")

            # 将截图转换为Base64字符串
            buffer = BytesIO()
            img.save(buffer, format="PNG")
            img_str = base64.b64encode(buffer.getvalue()).decode()
            # 发送Base64字符串给客户端
            await websocket.send(img_str)
        except Exception as e:
            print(f"Error handling screenshot request: {e}")

async def start_server():
    global server
    server = await websockets.serve(handle_screenshot, '127.0.0.1', 12345)
    print("Server started")

def start_or_stop_server():
    global server, loop
    if server is None:
        # 提交任务到事件循环
        asyncio.run_coroutine_threadsafe(start_server(), loop)
        button.config(text="停止服务")
        messagebox.showinfo("信息", "WebSocket服务已启动")
    else:
        server.close()
        server = None
        button.config(text="启动服务")
        messagebox.showinfo("信息", "WebSocket服务已停止")

def run_event_loop():
    global loop
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    loop.run_forever()

# 创建一个线程来运行事件循环
thread = threading.Thread(target=run_event_loop, daemon=True)
thread.start()

# 创建Tkinter窗口
root = tk.Tk()
root.title("WebSocket截图服务")
root.geometry("350x230")  # 设置窗口大小

# 添加文本
label = tk.Label(root, text="注意事项:\n1. 请求地址'ws://127.0.0.1:12345'。\n2. 传递四个参数x1, y1, x2, y2，\n分别是左上角坐标和右下角坐标。\njs语言示例：ws.send(`${x1},${y1},${x2},${y2}`);\n3. 返回base64字符串。")
label.pack(pady=20)

# 添加按钮
button = tk.Button(root, text="启动服务", command=start_or_stop_server)
button.pack(pady=20)

# 运行Tkinter主循环
root.mainloop()